package org.nicta.learn.factorization.GreedyResidualFitting;

import java.util.HashMap;
import java.util.HashSet;
import java.util.Random;

import org.nicta.colt.utils.ColtUtils;
import org.nicta.datasetLoader.MovieLens.MovieLensDataLoader;
import org.nicta.learn.optimization.LBFGS;

import cern.colt.matrix.DoubleFactory2D;
import cern.colt.matrix.DoubleMatrix2D;
import cern.colt.matrix.linalg.Algebra;
import cern.jet.math.Functions;

public class GRF {
	private int dimension;
	private String ratingsFileLocation;
	private MovieLensDataLoader mlLoader;

	DoubleMatrix2D U;
	DoubleMatrix2D V;
	Algebra algebra;
	
	public HashMap<Integer[], Double> ratings;
	public HashMap<Integer[], Double> testRatings;
	public HashMap<Integer, HashSet<Integer>> userRatedItems;
	public HashMap<Integer, HashSet<Integer>> itemsUsersMap;
	public HashMap<Integer, HashMap<Integer, Double>> itemUserRatings;
	public HashMap<Integer, Double> userAverage;

	public double lambdaUser;
	public double lambdaItem;
	
	public GRF(int dimension, String location,String datasettype) throws Exception{
		this.dimension = dimension;
		this.ratingsFileLocation = location;
		loadMovieLensData(datasettype);
		testRatings = createTestSet(ratings,0.1);
		userAverage = getUserAverage(itemUserRatings, userRatedItems);
	}
	
	public void loadMovieLensData(String datasetType) throws Exception{
		mlLoader = new MovieLensDataLoader(datasetType, ratingsFileLocation,
				null, null);
		mlLoader.loadData();
		ratings = mlLoader.getRatings();
		userRatedItems = mlLoader.getUserRatedMoviesMap();
		itemsUsersMap = mlLoader.getMoviesUsersMap();
		itemUserRatings = mlLoader.getmovieUserRatings();
		algebra = new Algebra();
		lambdaUser = 10;
		lambdaItem = 10;
	}
	
	public HashMap<Integer[], Double> createTestSet(
			HashMap<Integer[], Double> ratings, double fraction ){
		HashMap<Integer[], Double> testRatings;
		testRatings = new HashMap<Integer[], Double>();
		Object[] ratingsKey  = ratings.keySet().toArray();
		double testSetSize = fraction * ratings.keySet().size();
		Random rand = new Random();
		int key;
		while(testRatings.size() < testSetSize){
			key = rand.nextInt(ratingsKey.length);
			Object userItemPair = ratingsKey[key];
			//TODO:needs optimization by the use of HashSet
			if(!ratings.containsKey(userItemPair)) continue;
			Integer userId = ((Integer[]) userItemPair)[0];
			Integer movieId = ((Integer[]) userItemPair)[1];
			if(itemsUsersMap.get(movieId).size() > 1 && 
					userRatedItems.get(userId).size() > 1){
				Double rating =  ratings.get(userItemPair);
				ratings.remove(userItemPair);
				userRatedItems.get(userId).remove(movieId);
				itemsUsersMap.get(movieId).remove(userId);
				testRatings.put(new Integer[]{userId, movieId}, rating);
			}
		}
		return testRatings;
	}
	
	public HashMap<Integer[], Double> getResidue(){
		double rating;
		double predicted;
		if(U == null || V == null){
			return ratings;
		}
		HashMap<Integer[], Double> residue = new HashMap<Integer[], Double>();
        //Object[] ratingsKey  = ratings.keySet().toArray();
		for(Integer[] userItemPair: ratings.keySet()){
			Integer userId = userItemPair[0];
			Integer itemId = userItemPair[1];
			predicted = algebra.mult(U.viewColumn(userId-1),
					V.viewColumn(itemId-1));
			rating = ratings.get(userItemPair); 
			//System.out.format("Actual : %f Predicted :%f\n",rating,predicted);
			residue.put(new Integer[]{userId,itemId}, rating - predicted);
		}
		return residue;
	}
	
	public HashMap<Integer, Double> getUserAverage(
			HashMap<Integer, HashMap<Integer,Double>> itemUserRatings,
			HashMap<Integer, HashSet<Integer>> userRatedItems){
		double average;
		double sum;
		HashMap<Integer, Double> userAverage = new HashMap<Integer, Double>();
		for(Integer u: userRatedItems.keySet()){
			HashSet<Integer> uRatedMovies = userRatedItems.get(u);
			sum = 0;
			for(Integer m: uRatedMovies){
				sum += itemUserRatings.get(m).get(u);
			}
			average = sum/uRatedMovies.size();
			userAverage.put(u, average);
		}
		return userAverage;
	}
	
	public void ratingNormalizationPerUserAverage(
			HashMap<Integer[], Double> ratings,
			HashMap<Integer, Double> userAverage 
			){
		Object[] ratingsKey = ratings.keySet().toArray();
		for(Object userItemPair: ratingsKey){
			Double rating =  ratings.get(userItemPair);
			Integer userId = ((Integer[]) userItemPair)[0];
			Integer movieId = ((Integer[]) userItemPair)[1];
			rating = rating-userAverage.get(userId);
			ratings.put((Integer[]) userItemPair,rating);
			itemUserRatings.get(movieId).put(userId, rating);
		}
	}
	
	public DoubleMatrix2D getUserGradient(HashMap<Integer[], Double> residue,
			DoubleMatrix2D U, DoubleMatrix2D V){
		Integer userId;
		Integer itemId;
		double value;
		double predicted;
		double error;
		DoubleMatrix2D userGradient = DoubleFactory2D.dense.make(1,mlLoader.getUsersCount());
		for(Integer[] userItemPair: residue.keySet()){
			userId = userItemPair[0];
			itemId = userItemPair[1];
			value = residue.get(userItemPair);
			predicted = U.get(0, userId-1) * V.get(0, itemId-1);
			error = value - predicted;
			double temp = userGradient.get(0, userId-1);
			userGradient.set(0, userId-1, temp - error * V.get(0, itemId-1));
		}
		userGradient.assign(ColtUtils.multScalar(U, lambdaUser),Functions.plus);
		return userGradient;
	}
	
	public DoubleMatrix2D getMovieGradient(HashMap<Integer[], Double> residue,
			DoubleMatrix2D U, DoubleMatrix2D V){
		Integer userId;
		Integer itemId;
		double value;
		double predicted;
		double error;
		DoubleMatrix2D movieGradient = DoubleFactory2D.dense.make(1,mlLoader.getMoviesCount());
		for(Integer[] userItemPair: residue.keySet()){
			userId = userItemPair[0];
			itemId = userItemPair[1];
			value = residue.get(userItemPair);
			predicted = U.get(0, userId-1) * V.get(0, itemId-1);
			error = value - predicted;
			double temp = movieGradient.get(0, itemId-1);
			movieGradient.set(0, itemId-1, temp - error * U.get(0, userId-1));
		}
		movieGradient.assign(ColtUtils.multScalar(V, lambdaItem),Functions.plus);
		return movieGradient;
	}
	
	public double calculateCost(HashMap<Integer[], Double> residue,
			DoubleMatrix2D U, DoubleMatrix2D V){
		double cost = 0.0;
		double norm = 0.0;
		double predicted;
		int userId;
		int itemId;
		double value;
		for(Integer[] userItemPair: residue.keySet()){
			userId = userItemPair[0];
			itemId = userItemPair[1];
			value =  residue.get(userItemPair);
			predicted = U.get(0,userId-1) * V.get(0,itemId-1);
			cost += Math.pow((value-predicted), 2);
		}
		norm = (Math.pow(algebra.normF(U),2) * lambdaUser+ 
				Math.pow(algebra.normF(V),2)* lambdaItem)/2; 
		return cost/2+norm;
	}
	
	public double[] calculateTestError(){
		double rmseCost = 0.0;
		double maeCost = 0.0;
		double predicted;
		double value;
		int userId;
		int itemId;
		double[] result = new double [2];
		for(Integer[] userItemPair:testRatings.keySet()){
			userId = userItemPair[0];
			itemId = userItemPair[1];
			value = testRatings.get(userItemPair);
			predicted = algebra.mult(U.viewColumn(userId-1), V.viewColumn(itemId-1))
					+ userAverage.get(userId);
			//cutoff
			if (predicted > 5) predicted = 5;
			if (predicted < 1) predicted = 1;
			//System.out.format("Actual : %f Predicted: %f\n",value,predicted);
			rmseCost += Math.pow((value - predicted),2);
			maeCost += Math.abs(value - predicted);
		}
		result[0] = Math.sqrt(rmseCost*(1.0/testRatings.size()));
		result[1] = maeCost*(1.0/testRatings.size());
		return result; 
	}
	
	public DoubleMatrix2D[] rankOneApproximationGD(HashMap<Integer[], Double> residue)
			throws Exception{
		
		DoubleMatrix2D U1 = DoubleFactory2D.dense.random(1,mlLoader.getUsersCount()) ;
		DoubleMatrix2D V1 = DoubleFactory2D.dense.random(1,mlLoader.getMoviesCount()) ;
		double alpha = 0.001;
		DoubleMatrix2D userGradient;
		DoubleMatrix2D movieGradient;
		DoubleMatrix2D[] result = new DoubleMatrix2D[2];

		for(int i=0;i<500;i++){
			userGradient = getUserGradient(residue, U1, V1);
			movieGradient = getMovieGradient(residue, U1, V1);
			U1.assign(ColtUtils.multScalar(userGradient, alpha),Functions.minus);
			V1.assign(ColtUtils.multScalar(movieGradient, alpha),Functions.minus);
		}
		result[0] = U1;
		result[1] = V1;
		return result;
		
	}
	
	public DoubleMatrix2D[] rankOneApproximation(HashMap<Integer[], Double> residue)
			throws Exception{
		
		Random random = new Random();
		DoubleMatrix2D U1 = DoubleFactory2D.dense.random(1,mlLoader.getUsersCount()) ;
		DoubleMatrix2D V1 = DoubleFactory2D.dense.random(1,mlLoader.getMoviesCount()) ;
		
//		for(int i=0;i<mlLoader.getUsersCount();i++) {
//			U1.set(0, i, random.nextGaussian());
//		}
//		for(int i=0;i<mlLoader.getMoviesCount();i++) {
//			V1.set(0, i, random.nextGaussian());
//		}
		
		DoubleMatrix2D userGradient;
		DoubleMatrix2D movieGradient;
		DoubleMatrix2D[] result = new DoubleMatrix2D[2];
		double[] diag = new double[mlLoader.getMoviesCount() + 
	                                mlLoader.getUsersCount()];
		
		int[] iprint = {0,0};
		int[] iflag = {0};
		
		double[] variables = new double[mlLoader.getMoviesCount() + 
		                                mlLoader.getUsersCount()];
		
		double[] derivatives = new double[mlLoader.getMoviesCount() + 
			                                mlLoader.getUsersCount()];
		double error = Double.MAX_VALUE;
		double prev_error;
		
		do{
			prev_error = error;
			userGradient = getUserGradient(residue, U1, V1);
			movieGradient = getMovieGradient(residue, U1, V1);
			error = calculateCost(residue, U1, V1);
			ColtUtils.mergeFlattenTo1D(U1, V1, variables);
			ColtUtils.mergeFlattenTo1D(userGradient, movieGradient, derivatives);
			LBFGS.lbfgs(variables.length, 5, variables, error, derivatives,
					false, diag, iprint, 0.0001,
					1e-15, iflag);
			ColtUtils.assignByUnflattening(variables,U1,V1);	
		}while(Math.abs(error-prev_error) > 0.0001 && iflag[0] != 0);
		
		result[0] = U1;
		result[1] = V1;
		return result;
	}
	
	private DoubleMatrix2D appendColumn(DoubleMatrix2D U, DoubleMatrix2D U1){
		if(U == null) return U1;
		return DoubleFactory2D.dense.appendRows(U, U1);
	}
	
	public void runGrf() throws Exception{
		DoubleMatrix2D U1;
		DoubleMatrix2D V1;
		DoubleMatrix2D [] uvPair = new DoubleMatrix2D[2];
		HashMap<Integer[], Double> residue = null;
		for(int i = 0 ; i < dimension ; i++){
			residue = getResidue();
			double norm = 0 ;
			for(Integer[] userItemPair:residue.keySet()){
				norm += Math.abs(residue.get(userItemPair));
			}
			System.out.println("Residue : "+norm);
			uvPair = rankOneApproximation(residue);
			U = appendColumn(U, uvPair[0]);
			V = appendColumn(V, uvPair[1]);	
		}

	}
}
